[https://nvbugs/6506990][fix] Don't treat NVLE-only as CC enabled#16850
[https://nvbugs/6506990][fix] Don't treat NVLE-only as CC enabled#16850dhansen-nvidia wants to merge 1 commit into
Conversation
NVLE can be enabled independently of CC, so CC restrictions such as H2D bounce buffering do not apply to NVLE-only systems. Query and cache the CC and NVLE states together as one coherent NVML snapshot. Keep confidential_compute_enabled() as the existing boolean interface so current call sites remain unchanged, and cover modern and legacy NVML states with CPU-only tests. Signed-off-by: Samuel Mendoza-Jonas <smendozajona@nvidia.com> Signed-off-by: Dan Hansen <1+dhansen-nvidia@users.noreply.github.com>
WalkthroughAdds a cached NVML helper that reports confidential-compute and NVLink-encryption status, supports legacy fallback, and updates the existing boolean helper. New tests cover feature combinations, caching, delegation, and fallback behavior. ChangesConfidential compute status
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant get_cc_and_nvle_status
participant pynvml
Caller->>get_cc_and_nvle_status: request cached status
get_cc_and_nvle_status->>pynvml: query compute settings
pynvml-->>get_cc_and_nvle_status: return CC and NVLE settings
get_cc_and_nvle_status->>pynvml: attempt shutdown
get_cc_and_nvle_status-->>Caller: return status tuple
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_utils.py`:
- Around line 1212-1215: Update the CC/NVLE state query exception handling in
the surrounding function to catch only pynvml.NVMLError in both visible
handlers, allowing unexpected programming or API-compatibility exceptions to
propagate instead of returning (False, False).
In `@tests/unittest/utils/test_confidential_compute.py`:
- Around line 73-77: Add a test covering the ImportError path in
get_cc_and_nvle_status by making the pynvml import fail and asserting the result
is (False, False). Keep the existing cache-clearing fixture behavior, and
include an actionable coverage assessment for this test-code change.
- Around line 73-77: Annotate the clear_confidential_compute_status_cache
fixture and all referenced test functions with parameter and return type
annotations, using the appropriate fixture/test signatures and None for
functions that do not return a value. Apply this consistently to every function
in the affected test sections.
- Around line 117-124: Update test_confidential_compute_enabled to monkeypatch
tensorrt_llm._utils.get_cc_and_nvle_status with a tuple-returning stub and
assert confidential_compute_enabled() returns its first element, verifying
delegation and the single cached-query contract rather than independently
simulating NVML. Remove the direct pynvml setup from this test, and add an
actionable coverage assessment for the affected test-code change.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 95957461-80c5-4552-8057-85ec51e8fcc5
📒 Files selected for processing (2)
tensorrt_llm/_utils.pytests/unittest/utils/test_confidential_compute.py
| except Exception as e: | ||
| logger.error(f"Error querying confidential compute state: {str(e)}") | ||
| logger.error(f"Error querying CC and NVLE state: {str(e)}") | ||
| except Exception as e: | ||
| logger.error(f"Error querying confidential compute state: {str(e)}") | ||
| logger.error(f"Error querying CC and NVLE state: {str(e)}") |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Do not silently treat unexpected NVML/API failures as CC-disabled.
Catching Exception converts programming or API-compatibility failures into (False, False), which can bypass CC-specific restrictions. Catch pynvml.NVMLError here and let unexpected failures surface.
Proposed fix
- except Exception as e:
- logger.error(f"Error querying CC and NVLE state: {str(e)}")
- except Exception as e:
- logger.error(f"Error querying CC and NVLE state: {str(e)}")
+ except pynvml.NVMLError as error:
+ logger.error(f"Error querying CC and NVLE state: {error!s}")
+ except pynvml.NVMLError as error:
+ logger.error(f"Error querying CC and NVLE state: {error!s}")As per coding guidelines, “Catch the narrowest possible exceptions.”
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| except Exception as e: | |
| logger.error(f"Error querying confidential compute state: {str(e)}") | |
| logger.error(f"Error querying CC and NVLE state: {str(e)}") | |
| except Exception as e: | |
| logger.error(f"Error querying confidential compute state: {str(e)}") | |
| logger.error(f"Error querying CC and NVLE state: {str(e)}") | |
| except pynvml.NVMLError as error: | |
| logger.error(f"Error querying CC and NVLE state: {error!s}") | |
| except pynvml.NVMLError as error: | |
| logger.error(f"Error querying CC and NVLE state: {error!s}") |
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 1212-1212: Do not catch blind exception: Exception
(BLE001)
[warning] 1213-1213: Use explicit conversion flag
Replace with conversion flag
(RUF010)
[warning] 1214-1214: Do not catch blind exception: Exception
(BLE001)
[warning] 1215-1215: Use explicit conversion flag
Replace with conversion flag
(RUF010)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_utils.py` around lines 1212 - 1215, Update the CC/NVLE state
query exception handling in the surrounding function to catch only
pynvml.NVMLError in both visible handlers, allowing unexpected programming or
API-compatibility exceptions to propagate instead of returning (False, False).
Sources: Coding guidelines, Linters/SAST tools
| @pytest.fixture(autouse=True) | ||
| def clear_confidential_compute_status_cache(): | ||
| get_cc_and_nvle_status.cache_clear() | ||
| yield | ||
| get_cc_and_nvle_status.cache_clear() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Cover the missing-pynvml fallback.
The new ImportError branch in get_cc_and_nvle_status() is untested. Add a test that makes import pynvml fail and asserts (False, False).
As per path instructions, test-code changes require an actionable coverage assessment.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/utils/test_confidential_compute.py` around lines 73 - 77, Add
a test covering the ImportError path in get_cc_and_nvle_status by making the
pynvml import fail and asserting the result is (False, False). Keep the existing
cache-clearing fixture behavior, and include an actionable coverage assessment
for this test-code change.
Source: Path instructions
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add the required function annotations.
The autouse fixture and all test functions omit parameter and return annotations. As per coding guidelines, “Annotate every function.”
Also applies to: 89-96, 99-105, 117-124, 127-136
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/utils/test_confidential_compute.py` around lines 73 - 77,
Annotate the clear_confidential_compute_status_cache fixture and all referenced
test functions with parameter and return type annotations, using the appropriate
fixture/test signatures and None for functions that do not return a value. Apply
this consistently to every function in the affected test sections.
Source: Coding guidelines
| def test_confidential_compute_enabled(monkeypatch, cc_feature, multi_gpu_mode, expected): | ||
| pynvml = _make_pynvml( | ||
| cc_feature=cc_feature, | ||
| multi_gpu_mode=multi_gpu_mode, | ||
| ) | ||
| monkeypatch.setitem(sys.modules, "pynvml", pynvml) | ||
|
|
||
| assert confidential_compute_enabled() is expected |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Test delegation rather than only equivalent output.
This test would still pass if confidential_compute_enabled() independently queried NVML. Monkeypatch tensorrt_llm._utils.get_cc_and_nvle_status and assert the wrapper returns its first tuple element, preserving the single cached-query contract.
As per path instructions, test-code changes require an actionable coverage assessment.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/utils/test_confidential_compute.py` around lines 117 - 124,
Update test_confidential_compute_enabled to monkeypatch
tensorrt_llm._utils.get_cc_and_nvle_status with a tuple-returning stub and
assert confidential_compute_enabled() returns its first element, verifying
delegation and the single cached-query contract rather than independently
simulating NVML. Remove the direct pynvml setup from this test, and add an
actionable coverage assessment for the affected test-code change.
Source: Path instructions
1 participant
NVLE can be enabled independently of CC, so CC restrictions such as H2D bounce buffering do not apply to NVLE-only systems.
Query and cache the CC and NVLE states together as one coherent NVML snapshot. Keep confidential_compute_enabled() as the existing boolean interface so current call sites remain unchanged, and cover modern and legacy NVML states with CPU-only tests.
Dev Engineer Review
confidential_compute_enabled()as a boolean API while avoiding CC-only restrictions on NVLE-only systems.QA Engineer Review
test_get_cc_and_nvle_statustest_get_cc_and_nvle_status_is_cachedtest_confidential_compute_enabledtest_get_cc_and_nvle_status_uses_legacy_cc_fallbackDescription
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.